#include "plan_env/entropy_map.h"

void EntropyMap::initMap(ros::NodeHandle &nh)
{
  node_ = nh;
  this->GridMap::initMap(node_);

  node_.param("entropy_map/entropy_x_bound", emp_.entropy_x_bound_, 1);
  node_.param("entropy_map/entropy_y_bound", emp_.entropy_y_bound_, 1);
  node_.param("entropy_map/env_vel_max", emp_.env_vel_max_, 0.5);
  node_.param("entropy_map/entropy_inc_time", emp_.entropy_inc_time_, 6.0);
  node_.param("entropy_map/traj_attract_factor", emp_.traj_attract_factor_, 1.0);
  node_.param("entropy_map/safe_dis", emp_.safe_dis, 1.0);
  node_.param("entropy_map/dyn_obs_R2_max_", emp_.dyn_obs_R2_max_, 0.99);
  node_.param("entropy_map/lambda_dyn_obs_", emp_.lambda_dyn_obs_, 9.0);
  node_.param("entropy_map/dyn_attention_gain_", emp_.dyn_attention_gain_, 1.0);

  emp_.env_vel_max_voxel_ = emp_.env_vel_max_ * gmp_.resolution_inv_;

  int buffer_size = gmp_.map_voxel_num_(0) * gmp_.map_voxel_num_(1);
  emd_.entropy_time_buffer_ = vector<double>(buffer_size, -emp_.entropy_inc_time_);
  emd_.traj_buffer_ = vector<double>(buffer_size, 0.0);
  emd_.dyn_obs_buffer_ =
      vector<vector<double>>(emp_.max_dyn_buffer_size,
                             vector<double>(buffer_size, 0.0));

  // TODO MZY 还真不一定是这个，如果输入变了的话
  emd_.pts_bound_.resize((int)ceil((640 - gmp_.depth_filter_margin_ * 2) / gmp_.skip_pixel_));
  emd_.pts_bound_cnt_ = 0;

  entropy_timer_ = node_.createTimer(ros::Duration(0.05), &EntropyMap::updateEntropyCallback, this);
  map_entropy_pub_ = node_.advertise<sensor_msgs::PointCloud2>("entropy_map/entropy", 10);
  map_traj_info_pub_ = node_.advertise<sensor_msgs::PointCloud2>("entropy_map/traj_info", 10);
  map_dyn_obs_pub_ = node_.advertise<sensor_msgs::PointCloud2>("entropy_map/dyn_obs", 10);

  emd_.entropy_need_update_ = false;

  emd_.start_time_ = ros::Time::now();
}

void EntropyMap::updateEntropyCallback(const ros::TimerEvent &)
{
  if (!emd_.entropy_need_update_)
    return;

  if (emd_.pts_bound_cnt_ == 0)
    return;

  // ros::Time t1, t2, t3, t4;
  // t1 = ros::Time::now();

  /* get point on sensor range bound */
  Eigen::Matrix3d camera_r = gmd_.camera_r_m_;
  Eigen::Vector3d camera_p = gmd_.camera_pos_;
  int pts_cnt = emd_.pts_bound_cnt_;
  std::vector<Eigen::Vector3d> pts_bound(emd_.pts_bound_.begin(),
                                         emd_.pts_bound_.begin() + pts_cnt);
  std::vector<Eigen::Vector2d> pts_bound_2d;
  pts_bound_2d.resize(pts_cnt);
  Eigen::Vector3d proj_pt;
  for (int i = 0; i < pts_cnt; i++)
  {
    proj_pt = camera_r * pts_bound[i] + camera_p;
    if (!isInMap(proj_pt))
    {
      proj_pt = closetPointInMap(proj_pt, gmd_.camera_pos_);
    }
    pts_bound_2d[i] = Eigen::Vector2d(proj_pt(0), proj_pt(1));
  }
  // t2 = ros::Time::now();
  // TODO MZY 加上去除重复点的操作

  /* Raycasting and update entropy map*/
  ros::Time update_time = ros::Time::now();
  Eigen::Vector2i end_idx;
  posToIndex(Eigen::Vector2d(camera_p(0), camera_p(1)), end_idx);
  size_t bound_num = pts_bound_2d.size();
  for (size_t i = 0; i < bound_num; i++)
  {
    std::vector<Eigen::Vector2i> ray_points;
    Eigen::Vector2i bound_index;
    posToIndex(pts_bound_2d[i], bound_index);
    if (Raycast_2d(bound_index, end_idx, ray_points))
    {
      size_t ray_num = ray_points.size();
      for (size_t j = 0; j < ray_num; j++)
      {
        int id_ctns = toAddress(ray_points[j]);
        emd_.entropy_time_buffer_[id_ctns] =
            (update_time - emd_.start_time_).toSec();
      }
    }
  }
  // t3 = ros::Time::now();

  // cout << setprecision(7);
  // cout << "[ENTROPY]: t2=" << (t2 - t1).toSec()
  //      << " t3=" << (t3 - t2).toSec() << endl;
  // result [ENTROPY]: t2 = 4.015e-06 t3 = 0.000176316

  emd_.entropy_need_update_ = false;
}

/**
 * Project depth image to points in world frame
 * Store in proj_pts_
 *
 * @note override from CLASS GridMap
 */
void EntropyMap::projectDepthImage()
{
  if (!gmd_.has_first_depth_)
  {
    gmd_.has_first_depth_ = true;
    return;
  }

  /* used for 32FC1 */
  // const double inv_factor = 1.0 / gmp_.k_depth_scaling_factor_;
  /* used for 16UC1 */
  constexpr double inv_factor = 0.001;
  uint16_t *row_ptr;
  int cols = gmd_.depth_image_.cols;
  int rows = gmd_.depth_image_.rows;
  Eigen::Matrix3d camera_r = gmd_.camera_r_m_;
  Eigen::Vector3d camera_p = gmd_.camera_pos_;
  const int skip_pix = gmp_.skip_pixel_;
  const int margin = gmp_.depth_filter_margin_;
  double depth;
  Eigen::Vector3d pt_cur, pt_world;
  gmd_.proj_pts_cnt_ = 0;

  /* filter margin of depth image */
  for (int v = margin; v < rows - margin; v += skip_pix)
  {
    row_ptr = gmd_.depth_image_.ptr<uint16_t>(v) + margin;
    emd_.pts_bound_cnt_ = 0;
    for (int u = margin; u < cols - margin; u += skip_pix)
    {
      depth = (*row_ptr) * inv_factor;
      row_ptr = row_ptr + skip_pix;

      /* In case3, a special depth, will be processed in [raycastProcess] */
      if (depth < gmp_.depth_filter_mindist_)
      {
        continue;
      }
      else if (depth > gmp_.depth_filter_maxdist_)
      {
        depth = gmp_.max_ray_length_ + 0.1;
      }

      /* project to world frame */
      pt_cur(0) = (u - gmp_.cx_) * depth / gmp_.fx_;
      pt_cur(1) = (v - gmp_.cy_) * depth / gmp_.fy_;
      pt_cur(2) = depth;
      pt_world = camera_r * pt_cur + camera_p;
      gmd_.proj_pts_[gmd_.proj_pts_cnt_++] = pt_world;

      /* get point on sensor range bound */
      if (v == margin)
      {
        emd_.pts_bound_[emd_.pts_bound_cnt_++] = pt_cur;
      }
      else
      {
        if (emd_.pts_bound_[emd_.pts_bound_cnt_](2) > depth)
        {
          emd_.pts_bound_[emd_.pts_bound_cnt_] = pt_cur;
        }
        emd_.pts_bound_cnt_++;
      }
    }
  }
  emd_.entropy_need_update_ = true;
}

void EntropyMap::publishEntropy()
{
  if (map_entropy_pub_.getNumSubscribers() <= 0)
    return;

  int x_bound = emp_.entropy_x_bound_;
  int y_bound = emp_.entropy_y_bound_;

  pcl::PointXYZI pt;
  pcl::PointCloud<pcl::PointXYZI> cloud_entropy;
  Eigen::Vector2i odom_index;
  Eigen::Vector2d odom_pos{gmd_.camera_pos_(0), gmd_.camera_pos_(1)};
  posToIndex(odom_pos, odom_index);

  double cur_time = (ros::Time::now() - emd_.start_time_).toSec();

  for (int x = odom_index(0) - x_bound; x <= odom_index(0) + x_bound; ++x)
    for (int y = odom_index(1) - y_bound; y <= odom_index(1) + y_bound; ++y)
    {
      Eigen::Vector2i id{x, y};
      if (!isInMap(id)) continue;

      double value = calEnvEntropy(cur_time, toAddress(id));
      if (value == 1.0) continue;

      Eigen::Vector2d pos;
      indexToPos(id, pos);
      pt.x = pos(0);
      pt.y = pos(1);
      pt.z = 0.4;
      pt.intensity = value;
      cloud_entropy.push_back(pt);
    }

  cloud_entropy.width = cloud_entropy.points.size();
  cloud_entropy.height = 1;
  cloud_entropy.is_dense = true;
  cloud_entropy.header.frame_id = gmp_.frame_id_;
  sensor_msgs::PointCloud2 cloud_msg;
  pcl::toROSMsg(cloud_entropy, cloud_msg);
  map_entropy_pub_.publish(cloud_msg);
}

void EntropyMap::publishTrajMap()
{
  if (map_traj_info_pub_.getNumSubscribers() <= 0)
    return;

  int x_bound = emp_.entropy_x_bound_;
  int y_bound = emp_.entropy_y_bound_;

  pcl::PointXYZI pt;
  pcl::PointCloud<pcl::PointXYZI> cloud_entropy;
  Eigen::Vector2i odom_index;
  Eigen::Vector2d odom_pos{gmd_.camera_pos_(0), gmd_.camera_pos_(1)};
  posToIndex(odom_pos, odom_index);
  for (int x = odom_index(0) - x_bound; x <= odom_index(0) + x_bound; ++x)
    for (int y = odom_index(1) - y_bound; y <= odom_index(1) + y_bound; ++y)
    {
      Eigen::Vector2i id{x, y};
      if (!isInMap(id))
        continue;

      double value = emd_.traj_buffer_[toAddress(id)];
      if (value == 0.0)
        continue;

      Eigen::Vector2d pos;
      indexToPos(id, pos);
      pt.x = pos(0);
      pt.y = pos(1);
      pt.z = 0.6;
      pt.intensity = value;
      cloud_entropy.push_back(pt);
    }

  cloud_entropy.width = cloud_entropy.points.size();
  cloud_entropy.height = 1;
  cloud_entropy.is_dense = true;
  cloud_entropy.header.frame_id = gmp_.frame_id_;
  sensor_msgs::PointCloud2 cloud_msg;
  pcl::toROSMsg(cloud_entropy, cloud_msg);
  map_traj_info_pub_.publish(cloud_msg);
}

void EntropyMap::publishDynamicObs(bool isEmpty /* = false */)
{
  if (map_dyn_obs_pub_.getNumSubscribers() <= 0)
    return;

  pcl::PointCloud<pcl::PointXYZI> cloud_entropy;

  if (!isEmpty) {
    int x_bound = emp_.entropy_x_bound_;
    int y_bound = emp_.entropy_y_bound_;
    pcl::PointXYZI pt;
    Eigen::Vector2i odom_index;
    Eigen::Vector2d odom_pos{gmd_.camera_pos_(0), gmd_.camera_pos_(1)};
    posToIndex(odom_pos, odom_index);
    for (int x = odom_index(0) - x_bound; x <= odom_index(0) + x_bound; ++x)
      for (int y = odom_index(1) - y_bound; y <= odom_index(1) + y_bound; ++y)
      {
        Eigen::Vector2i id{x, y};
        if (!isInMap(id))
          continue;

        /* publish all layers */
        // for (size_t k = 0; k < emd_.dyn_layers_; k++) {
        //   double value = emd_.dyn_obs_buffer_[k][toAddress(id)];
        //   if (value == 0.0)
        //     continue;

        //   Eigen::Vector2d pos;
        //   indexToPos(id, pos);
        //   pt.x = pos(0);
        //   pt.y = pos(1);
        //   pt.z = 0.6 + 0.2 * k;
        //   pt.intensity = value;
        //   cloud_entropy.push_back(pt);
        // }

        /* publish the last layer */
        double value = emd_.dyn_obs_buffer_[0][toAddress(id)];
        if (value == 0.0)
          continue;

        Eigen::Vector2d pos;
        indexToPos(id, pos);
        pt.x = pos(0);
        pt.y = pos(1);
        pt.z = 0.8;
        pt.intensity = value;
        cloud_entropy.push_back(pt);
      }
  }

  cloud_entropy.width = cloud_entropy.points.size();
  cloud_entropy.height = 1;
  cloud_entropy.is_dense = true;
  cloud_entropy.header.frame_id = gmp_.frame_id_;
  sensor_msgs::PointCloud2 cloud_msg;
  pcl::toROSMsg(cloud_entropy, cloud_msg);
  map_dyn_obs_pub_.publish(cloud_msg);
}

void EntropyMap::visCallback(const ros::TimerEvent & /*event*/)
{
  if (!gmd_.has_odom_)
    return;

  publishMapInflate(true);
  publishMap();
  publishESDF();
  publishEntropy();
}

/**
 * @brief Get the ENTROPY started at [pos], the angle start from
 *         [theta], add by [delta]
 */
bool EntropyMap::getEntroy(Eigen::Vector3d pos, double theta, double delta,
                           int step_cnt, std::vector<double> &result)
{
  if (!isInMap(pos))
  {
    cerr << "[EntropyMap]: uav is not in map" << endl;
    return false;
  }

  result.resize(step_cnt, 0);

  preProcess(ros::Time::now().toSec());

  double cur_time = (ros::Time::now() - emd_.start_time_).toSec();

  bool isBlock = false;
  Eigen::Vector2i block_idx;

  for (int i = 0; i < step_cnt; i++)
  {
    result[i] = getEntroyAlongLine(cur_time, pos, theta, isBlock, block_idx);
    theta += delta;
  }

  return true;
}

/**
 * @brief
 *
 * @param pos
 * @param vel
 * @param t_plan
 * @param angles angles in world frames
 * @param result
 * @return true
 * @return false
 */
bool EntropyMap::getEntroy(Eigen::Vector3d pos, Eigen::Vector3d vel,
                           ros::Time t_plan,
                           std::vector<double> angles,
                           std::vector<double> &result)
{
  if (!isInMap(pos))
  {
    cerr << "[EntropyMap]: uav is not in map" << endl;
    return false;
  }

  preProcess(t_plan.toSec());

  size_t n = angles.size();
  result.resize(n, 0);

  double delta = 1.0;
  if (n > 1) delta = angles[1] - angles[0];

  double cur_time = (t_plan - emd_.start_time_).toSec();

  bool isBlock = false;
  bool isBlock_last = false;
  Eigen::Vector2i block_idx;
  Eigen::Vector2i block_idx_last;
  Eigen::Vector2d block_pos;
  std::vector<Eigen::Vector2d> block_pos_list;

  for (size_t i = 0; i < n; i++)
  {
    result[i] = getEntroyAlongLine(cur_time, pos, angles[i], isBlock, block_idx);

    if (isBlock && !isBlock_last)
    {
      indexToPos(block_idx, block_pos);
      block_pos_list.push_back(block_pos);
    }
    else if (!isBlock && isBlock_last)
    {
      indexToPos(block_idx_last, block_pos);
      block_pos_list.push_back(block_pos);
    }
    isBlock_last = isBlock;
    block_idx_last = block_idx;
  }

  pi_.reset(); // Must add this line

  if (block_pos_list.size() > 1)
  {
    double n_v = atan2(vel(1), vel(0));
    if (delta > 0)
      n_v += M_PI_2;
    else
      n_v -= M_PI_2;

    Eigen::Vector2d bias_p{gmp_.resolution_ * cos(n_v),
                           gmp_.resolution_ * sin(n_v)};

    block_pos_list.push_back(pos.head(2) + bias_p);
    block_pos_list.push_back(pos.head(2) - bias_p);
    pi_.addPolygon(block_pos_list);
  }

  return true;
}

/**
 * @brief Get the ENTROPY started at [pos] towards [theta],
 *        until [max_ray_length_]
 * @param tehta ray `s orientation in world frame
 *
 * @note must call [preProcess] before this function
 */
double EntropyMap::getEntroyAlongLine(
    double cur_time, Eigen::Vector3d pos, double theta,
    bool &isBlock, Eigen::Vector2i &block_idx)
{
  isBlock = false;
  std::vector<Eigen::Vector2i> ray_pts;
  Eigen::Vector3i start_idx, end_idx;
  // TODO 1.3来自1/cos(fov/2)
  // 20240125: 删除了1.3，认为“期望范围”应该小于“实际观测范围”，
  //           不然会一直看向看不到的地方
  Eigen::Vector3d end_pt{pos(0) + gmp_.max_ray_length_ * cos(theta),
                         pos(1) + gmp_.max_ray_length_ * sin(theta),
                         pos(2)};
  if (!isInMap(end_pt))
  {
    end_pt = closetPointInMap(end_pt, gmd_.camera_pos_);
  }
  posToIndex(pos, start_idx);
  posToIndex(end_pt, end_idx);
  Raycast_2d(start_idx.head(2), end_idx.head(2), ray_pts);

  double result = 0;
  size_t ray_num = ray_pts.size();
  for (size_t j = 0; j < ray_num; j++)
  {
    Eigen::Vector2i idx = ray_pts[j];
    Eigen::Vector3i idx_3{idx(0), idx(1), start_idx(2)};

    if (getOccupancy_idx(idx_3)) {
      isBlock = true;
      block_idx = idx;
      break;
    }

    int address = toAddress(idx);
    // TODO result怎么加还是非常关键的 最重要的
    double f_entropy = calEnvEntropy(cur_time, address);
    double f_dyn_obs = calDynamicObsEntropy(address);
    double f_esdf = calDistanceCost(idx_3);
    double f_traj = emd_.traj_buffer_[address];
    result += (f_entropy + emp_.lambda_dyn_obs_ * f_dyn_obs) *
              (f_traj + f_esdf + 1);
  }
  return result;
}

/**
 * @brief Fuse the traj information into map,
 *        fill a circle area orign at [pos_3], radius = vel * [t_cur-t_start]
 * @param t_end the duration of traj
 *
 * TODO 感觉有点消耗资源
 */
void EntropyMap::updateTrajMap(std::vector<Eigen::Vector3d> pos_3,
                               std::vector<double> t_cur,
                               double duration,
                               double start_time_global)
{
  if (pos_3.empty()) return;
  if (pos_3.size() != t_cur.size())
  {
    cout << RED_B << "[EntropyMap::updateTrajMap]: input error, bug"
         << RESET << endl;
    return;
  }

  emd_.traj_update_time_ = start_time_global;
  emd_.traj_duration_ = duration;

  size_t n = pos_3.size();
  for (size_t i = 0; i < n; i++)
  {
    // TODO value 具体怎么填充有待商榷 还有这个radius
    double value = 1 - t_cur[i] / duration;
    int radius = (int)floor(emp_.env_vel_max_voxel_ * t_cur[i]);
    int radius_square = radius * radius;
    Eigen::Vector2i origin, point;
    posToIndex(Eigen::Vector2d(pos_3[i](0), pos_3[i](1)), origin);
    for (int y = -radius; y <= radius; ++y)
    {
      for (int x = -radius; x <= radius; ++x)
      {
        if (x * x + y * y > radius_square)
          continue;

        point << origin(0) + x, origin(1) + y;
        if (!isInMap(point))
          continue;

        if (emd_.traj_buffer_[toAddress(point)] >= value)
          continue;

        emd_.traj_buffer_[toAddress(point)] = value;
      }
    }
  }
}

/**
 * @brief update entropy map with dynamic obstacles
 *
 * @param pos   :dynamic obs position in future
 * @param R2    :goodness of fitness of the dynamic obs corresponding to pos
 * @param times :global time of the dynamic obs in future corresponding to pos
 * @param isdanger :whether the dynamic obs will collide with uav
 */
void EntropyMap::updateDynamicObs(
    const std::vector<std::vector<Eigen::Vector3d>> &pos,
    const std::vector<std::vector<double>> &R2,
    const std::vector<std::vector<double>> &times,
    const std::vector<bool> &isdanger)
{
  if (pos.size() != R2.size() ||
      pos.size() != isdanger.size() ||
      pos.size() != times.size())
  {
    cout << RED_B << "[EntropyMap::updateDynamicObs]: input error, bug"
         << RESET << endl;
    return;
  }

  if (pos.empty()) return;

  int radius = 3 * gmp_.resolution_inv_;
  int radius_square = radius * radius;

  emd_.dyn_layers_ = 0;
  size_t index = 0;
  for (size_t i = 0; i < pos.size(); i++)
  {
    if (pos[i].size() > emd_.dyn_layers_) {
      emd_.dyn_layers_ = pos[i].size();
      index = i;
    }
  }
  emd_.dyn_layers_ = min(emd_.dyn_layers_, emp_.max_dyn_buffer_size);
  emd_.dyn_time_buffer_.clear();
  emd_.dyn_time_buffer_.insert(emd_.dyn_time_buffer_.begin(),
                               times[index].begin(),
                               times[index].begin() + emd_.dyn_layers_);

  for (size_t i = 0; i < pos.size(); i++)
  {
    double danger_value = 0.0;
    if (isdanger[i]) danger_value = emp_.dyn_attention_gain_;

    for (size_t k = 0; k < emd_.dyn_layers_; k++)
    {
      size_t j = k;
      if (j >= pos[i].size()) j = pos[i].size() - 1;

      // cout << "[EntropyMap::updateDynamicObs]: R2  = " << R2 << endl;

      double delta = emp_.dyn_obs_R2_max_ - R2[i][j];
      if (delta <= 0 && !isdanger[i]) continue;

      // TODO value 具体怎么填充有待商榷
      double value = danger_value;
      if (delta > 0) value += 2.0 / (1 + exp(-500 * delta)) - 1.0;

      // cout << GREEN_B << "isDanger: " << isdanger[i] << endl;
      // cout << "value  = " << value << RESET << endl;

      Eigen::Vector2i origin, point;
      posToIndex(Eigen::Vector2d(pos[i][j](0), pos[i][j](1)), origin);
      for (int y = -radius; y <= radius; ++y)
      {
        for (int x = -radius; x <= radius; ++x)
        {
          int distance = x * x + y * y;
          if (distance > radius_square)
            continue;

          point << origin(0) + x, origin(1) + y;
          if (!isInMap(point))
            continue;

          int addr = toAddress(point);

          emd_.dyn_obs_buffer_[k][addr] =
              max(emd_.dyn_obs_buffer_[k][addr],
                  value * (1 - double(distance) / double(radius_square)));
        }
      }
    }
  }
}

/**
 * @brief reset traj map buffer to zero before filling the value
 */
void EntropyMap::resetTrajMap()
{
  std::fill(emd_.traj_buffer_.begin(), emd_.traj_buffer_.end(), 0.0);
}

void EntropyMap::resetDynamicObs()
{
  for (size_t i = 0; i < emd_.dyn_layers_; i++) {
    std::fill(emd_.dyn_obs_buffer_[i].begin(),
              emd_.dyn_obs_buffer_[i].end(), 0.0);
  }
  emd_.dyn_time_buffer_.clear();
  emd_.dyn_layers_ = 0;
}